#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import random
import json


def gen_json(inpath: str, outpath: str):
    base_json = {"source_org": "SINH", "disease_name": "COVID-19", "batch_id": "covid19-001"}
    triple_list = []
    infile = open(inpath, "r")
    infile.readline()
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        line_ht = {}
        line_ht["source_entity_type"] = splitline[2]
        line_ht["source_entity_id"] = "covid" + str(random.randrange(0, 10000))
        if splitline[1].startswith('"') and splitline[1].endswith('"'):
            line_ht["source_entity_name"] = splitline[1][1:-1]
        else:
            line_ht["source_entity_name"] = splitline[1]
        line_ht["entity_relation_name"] = splitline[5]
        line_ht["target_entity_type"] = splitline[4]
        line_ht["target_entity_id"] = "covid" + str(random.randrange(0, 10000))
        if splitline[3].startswith('"') and splitline[3].endswith('"'):
            line_ht["target_entity_name"] = splitline[3][1:-1]
        else:
            line_ht["target_entity_name"] = splitline[3]
        triple_list.append(line_ht)
    infile.close()
    base_json["triple_list"] = triple_list
    outfile = open(outpath, "w")
    outfile.write(json.dumps(base_json) + "\n")
    outfile.flush()
    outfile.close()


if __name__ == '__main__':
    gen_json("D:/sjy_triple_data.txt", "D:/sjy_triple_data.json")
